package repository

import (
	"context"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"techoiceness.com/aiagent/llm-gateway/internal/entity"
)

// GetToolDatabaseSchema 获取工具数据库的架构信息
func (r *schemaRepository) GetToolDatabaseSchema(ctx context.Context, toolDBConnection *entity.ToolDBConnection) ([]TableInfo, error) {
	// 创建工具数据库连接
	toolDB, err := gorm.Open(postgres.Open(toolDBConnection.GetDSN()), &gorm.Config{})
	if err != nil {
		return nil, err
	}

	// 获取底层的sql.DB以便关闭连接
	sqlDB, err := toolDB.DB()
	if err != nil {
		return nil, err
	}
	defer sqlDB.Close()

	var columns []TableColumn

	query := `
	SELECT 
		n.nspname as schema_name,
		c.relname as table_name,
		a.attname as column_name,
		a.attnum as column_position,
		pg_catalog.format_type(a.atttypid, a.atttypmod) as data_type,
		a.attnotnull as not_null,
		a.atthasdef as has_default,
		pg_get_expr(ad.adbin, ad.adrelid) as default_value,
		col_description(a.attrelid, a.attnum) as column_comment,
		obj_description(c.oid) as table_comment
	FROM 
		pg_catalog.pg_class c
	JOIN 
		pg_catalog.pg_namespace n ON n.oid = c.relnamespace
	JOIN 
		pg_catalog.pg_attribute a ON a.attrelid = c.oid
	LEFT JOIN 
		pg_catalog.pg_attrdef ad ON ad.adrelid = c.oid AND ad.adnum = a.attnum
	WHERE 
		c.relkind = 'r'  -- 只查询表
		AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast')
		AND a.attnum > 0  -- 排除系统列
		AND NOT a.attisdropped  -- 排除已删除的列
	ORDER BY 
		n.nspname, c.relname, a.attnum
	`

	if err := toolDB.WithContext(ctx).Raw(query).Scan(&columns).Error; err != nil {
		return nil, err
	}

	// 将列信息按表分组
	tableMap := make(map[string]*TableInfo)
	for _, col := range columns {
		tableKey := col.SchemaName + "." + col.TableName

		if table, exists := tableMap[tableKey]; exists {
			table.Columns = append(table.Columns, col)
		} else {
			tableMap[tableKey] = &TableInfo{
				SchemaName:   col.SchemaName,
				TableName:    col.TableName,
				TableComment: col.TableComment,
				Columns:      []TableColumn{col},
			}
		}
	}

	// 转换为数组
	var tables []TableInfo
	for _, table := range tableMap {
		tables = append(tables, *table)
	}

	return tables, nil
}

// GetToolDatabaseTableList 获取工具数据库的表列表
func (r *schemaRepository) GetToolDatabaseTableList(ctx context.Context, toolDBConnection *entity.ToolDBConnection) ([]TableInfo, error) {
	// 创建工具数据库连接
	toolDB, err := gorm.Open(postgres.Open(toolDBConnection.GetDSN()), &gorm.Config{})
	if err != nil {
		return nil, err
	}

	// 获取底层的sql.DB以便关闭连接
	sqlDB, err := toolDB.DB()
	if err != nil {
		return nil, err
	}
	defer sqlDB.Close()

	var tables []TableInfo

	query := `
	SELECT DISTINCT
		n.nspname as schema_name,
		c.relname as table_name,
		obj_description(c.oid) as table_comment
	FROM 
		pg_catalog.pg_class c
	JOIN 
		pg_catalog.pg_namespace n ON n.oid = c.relnamespace
	WHERE 
		c.relkind = 'r'  -- 只查询表
		AND n.nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast')
	ORDER BY 
		n.nspname, c.relname
	`

	if err := toolDB.WithContext(ctx).Raw(query).Scan(&tables).Error; err != nil {
		return nil, err
	}

	return tables, nil
}
